// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/decoded_args.h>
#include <pollux/expression/vector_function.h>
#include <pollux/type/filter.h>

namespace kumo::pollux::functions {
    namespace {
        // This implements InPredicate using a set over VectorValues (pairs of
        // BaseVector, index). Can be used in place of Filters for Types not supported
        // by Filters or when custom comparisons are needed.
        // Returns NULL if
        // - input value is NULL
        // - in-list is NULL or empty
        // - input value doesn't have an exact match, but has an indeterminate match in
        // the in-list. E.g., 'array[null] in (array[1])' or 'array[1] in
        // (array[null])'.
        class VectorSetInPredicate : public exec::VectorFunction {
        public:
            struct VectorValue {
                BaseVector *vector;
                vector_size_t index;
            };

            struct VectorValueHash {
                size_t operator()(VectorValue value) const {
                    return value.vector->hash_value_at(value.index);
                }
            };

            struct VectorValueEqualTo {
                bool operator()(VectorValue left, VectorValue right) const {
                    return left.vector->equal_value_at(right.vector, left.index, right.index);
                }
            };

            using VectorSet =
            melon::F14FastSet<VectorValue, VectorValueHash, VectorValueEqualTo>;

            VectorSetInPredicate(
                VectorSet uniqueValues,
                bool hasNull,
                VectorPtr originalValues)
                : uniqueValues_{std::move(uniqueValues)},
                  hasNull_{hasNull},
                  originalValues_{std::move(originalValues)} {
            }

            static std::shared_ptr<exec::VectorFunction>
            create(const VectorPtr &values, vector_size_t offset, vector_size_t size) {
                VectorSet uniqueValues;
                bool hasNull = false;

                for (auto i = offset; i < offset + size; i++) {
                    if (values->contains_null_at(i)) {
                        hasNull = true;
                    }
                    uniqueValues.insert({values.get(), i});
                }

                return std::make_shared<VectorSetInPredicate>(
                    std::move(uniqueValues), hasNull, values);
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                const auto &arg = args[0];

                context.ensure_writable(rows, BOOLEAN(), result);
                result->clear_nulls(rows);
                auto *boolResult = result->as_unchecked<FlatVector<bool> >();

                context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
                    if (arg->is_null_at(row)) {
                        boolResult->set_null(row, true);
                    } else {
                        const bool found =
                                uniqueValues_.contains(VectorValue({arg.get(), row}));
                        if (found) {
                            if (arg->contains_null_at(row)) {
                                boolResult->set_null(row, true);
                            } else {
                                boolResult->set(row, true);
                            }
                        } else {
                            if ((arg->contains_null_at(row) || hasNull_) &&
                                hasIndeterminateMatch(arg, row)) {
                                boolResult->set_null(row, true);
                            } else {
                                boolResult->set(row, false);
                            }
                        }
                    }
                });
            }

        private:
            bool hasIndeterminateMatch(const VectorPtr &vector, vector_size_t index)
            const {
                for (const auto &value: uniqueValues_) {
                    if (!vector
                        ->equal_value_at(
                            value.vector,
                            index,
                            value.index,
                            CompareFlags::NullHandlingMode::kNullAsIndeterminate)
                        .has_value()) {
                        return true;
                    }
                }
                return false;
            }

            // Set of unique values to check against. This set doesn't include any value
            // that is null or contains null.
            const VectorSet uniqueValues_;

            // Boolean indicating whether one of the value was null or contained null.
            const bool hasNull_;

            // Vector being referenced by the 'uniqueValues_'.
            const VectorPtr originalValues_;
        };

        // Read 'size' values from 'valuesVector' starting at 'offset', de-duplicate
        // remove nulls and sort. Return a list of unique non-null values sorted in
        // ascending order and a boolean indicating whether there were any null values.
        template<typename T, typename U = T>
        std::pair<std::vector<T>, bool> toValues(
            const VectorPtr &valuesVector,
            vector_size_t offset,
            vector_size_t size) {
            auto simpleValues = valuesVector->as<SimpleVector<U> >();

            bool nullAllowed = false;
            std::vector<T> values;
            values.reserve(size);

            for (auto i = offset; i < offset + size; i++) {
                if (simpleValues->is_null_at(i)) {
                    nullAllowed = true;
                } else {
                    if constexpr (std::is_same_v<U, Timestamp>) {
                        values.emplace_back(simpleValues->value_at(i).toMillis());
                    } else {
                        values.emplace_back(simpleValues->value_at(i));
                    }
                }
            }

            // In-place sort, remove duplicates, and later std::move to save memory
            std::sort(values.begin(), values.end());
            auto last = std::unique(values.begin(), values.end());
            values.resize(std::distance(values.begin(), last));

            return {std::move(values), nullAllowed};
        }

        // Creates a filter for constant values. A null filter means either
        // no values or only null values. The boolean is true if the list is
        // non-empty and consists of nulls only.
        template<typename T>
        std::pair<std::unique_ptr<common::Filter>, bool> createBigintValuesFilter(
            const VectorPtr &valuesVector,
            vector_size_t offset,
            vector_size_t size) {
            auto valuesPair = toValues<int64_t, T>(valuesVector, offset, size);

            const auto &values = valuesPair.first;
            bool nullAllowed = valuesPair.second;

            if (values.empty() && nullAllowed) {
                return {nullptr, true};
            }
            POLLUX_USER_CHECK(
                !values.empty(),
                "IN predicate expects at least one non-null value in the in-list");
            if (values.size() == 1) {
                return {
                    std::make_unique<common::BigintRange>(
                        values[0], values[0], nullAllowed),
                    false
                };
            }

            return {common::createBigintValues(values, nullAllowed), false};
        }

        // For double, cast double to Int64 and reuse Int64 filters
        // For float, cast float to Int32 and promote to Int64
        template<typename T>
        std::pair<std::unique_ptr<common::Filter>, bool>
        createFloatingPointValuesFilter(
            const VectorPtr &valuesVector,
            vector_size_t offset,
            vector_size_t size) {
            auto valuesPair = toValues<T, T>(valuesVector, offset, size);

            auto &values = valuesPair.first;
            bool nullAllowed = valuesPair.second;

            if (values.empty() && nullAllowed) {
                return {nullptr, true};
            }
            POLLUX_USER_CHECK(
                !values.empty(),
                "IN predicate expects at least one non-null value in the in-list");
            // Avoid using FloatingPointRange for optimization of a single value in-list
            // as it does not support NaN as a bound for specifying a range.
            std::vector<int64_t> intValues(values.size());
            for (size_t i = 0; i < values.size(); ++i) {
                if (std::isnan(values[i])) {
                    // We de-normalize NaN values to ensure different binary representations
                    // are treated the same.
                    values[i] = std::numeric_limits<T>::quiet_NaN();
                }
                if constexpr (std::is_same_v<T, float>) {
                    if (values[i] == float{}) {
                        values[i] = 0;
                    }
                    intValues[i] = reinterpret_cast<const int32_t &>(
                        values[i]); // silently promote to int64
                } else {
                    if (values[i] == double{}) {
                        values[i] = 0;
                    }
                    intValues[i] = reinterpret_cast<const int64_t &>(values[i]);
                }
            }
            return {common::createBigintValues(intValues, nullAllowed), false};
        }

        // See createBigintValuesFilter.
        template<typename T>
        std::pair<std::unique_ptr<common::Filter>, bool> createHugeintValuesFilter(
            const VectorPtr &valuesVector,
            vector_size_t offset,
            vector_size_t size) {
            auto valuesPair = toValues<int128_t, T>(valuesVector, offset, size);

            const auto &values = valuesPair.first;
            bool nullAllowed = valuesPair.second;

            if (values.empty() && nullAllowed) {
                return {nullptr, true};
            }
            POLLUX_USER_CHECK(
                !values.empty(),
                "IN predicate expects at least one non-null value in the in-list");
            if (values.size() == 1) {
                return {
                    std::make_unique<common::HugeintRange>(
                        values[0], values[0], nullAllowed),
                    false
                };
            }

            return {common::createHugeintValues(values, nullAllowed), false};
        }

        // See createBigintValuesFilter.
        std::pair<std::unique_ptr<common::Filter>, bool> createBytesValuesFilter(
            const VectorPtr &valuesVector,
            vector_size_t offset,
            vector_size_t size) {
            auto valuesPair =
                    toValues<std::string, StringView>(valuesVector, offset, size);

            const auto &values = valuesPair.first;
            bool nullAllowed = valuesPair.second;
            if (values.empty() && nullAllowed) {
                return {nullptr, true};
            }

            POLLUX_USER_CHECK(
                !values.empty(),
                "IN predicate expects at least one value in the in-list");
            if (values.size() == 1) {
                return {
                    std::make_unique<common::BytesRange>(
                        values[0], false, false, values[0], false, false, nullAllowed),
                    false
                };
            }

            return {std::make_unique<common::BytesValues>(values, nullAllowed), false};
        }

        /// x IN (2, null) returns null when x != 2 and true when x == 2.
        /// Null for x always produces null, regardless of 'IN' list.
        class InPredicate : public exec::VectorFunction {
        public:
            explicit InPredicate(std::unique_ptr<common::Filter> filter, bool alwaysNull)
                : filter_{std::move(filter)}, alwaysNull_(alwaysNull) {
            }

            static std::shared_ptr<exec::VectorFunction> create(
                const std::string & /*name*/,
                const std::vector<exec::VectorFunctionArg> &inputArgs,
                const core::QueryConfig & /*config*/) {
                POLLUX_CHECK_EQ(inputArgs.size(), 2);
                auto inListType = inputArgs[1].type;

                POLLUX_CHECK_EQ(inListType->kind(), TypeKind::ARRAY);

                const auto &values = inputArgs[1].constantValue;
                POLLUX_USER_CHECK_NOT_NULL(
                    values, "IN predicate supports only constant IN list");

                if (values->is_null_at(0)) {
                    return std::make_shared<InPredicate>(nullptr, true);
                }

                POLLUX_CHECK_EQ(values->type_kind(), TypeKind::ARRAY);

                auto constantInput =
                        std::dynamic_pointer_cast<ConstantVector<ComplexType> >(values);

                auto array_vector = dynamic_cast<const ArrayVector *>(
                    constantInput->value_vector()->wrapped_vector());
                vector_size_t arrayIndex =
                        constantInput->value_vector()->wrapped_index(constantInput->index());
                auto offset = array_vector->offsetAt(arrayIndex);
                auto size = array_vector->sizeAt(arrayIndex);
                try {
                    POLLUX_USER_CHECK_GT(size, 0, "IN list must not be empty");
                } catch (...) {
                    return std::make_shared<exec::AlwaysFailingVectorFunction>(
                        std::current_exception());
                }

                const auto &elements = array_vector->elements();
                const auto &elementType = elements->type();

                if (elementType->providesCustomComparison()) {
                    return VectorSetInPredicate::create(elements, offset, size);
                }

                std::pair<std::unique_ptr<common::Filter>, bool> filter;

                switch (elementType->kind()) {
                    case TypeKind::HUGEINT:
                        filter = createHugeintValuesFilter<int128_t>(elements, offset, size);
                        break;
                    case TypeKind::BIGINT:
                        filter = createBigintValuesFilter<int64_t>(elements, offset, size);
                        break;
                    case TypeKind::INTEGER:
                        filter = createBigintValuesFilter<int32_t>(elements, offset, size);
                        break;
                    case TypeKind::SMALLINT:
                        filter = createBigintValuesFilter<int16_t>(elements, offset, size);
                        break;
                    case TypeKind::TINYINT:
                        filter = createBigintValuesFilter<int8_t>(elements, offset, size);
                        break;
                    case TypeKind::REAL:
                        filter = createFloatingPointValuesFilter<float>(elements, offset, size);
                        break;
                    case TypeKind::DOUBLE:
                        filter =
                                createFloatingPointValuesFilter<double>(elements, offset, size);
                        break;
                    case TypeKind::BOOLEAN:
                        // Hack: using BIGINT filter for bool, which is essentially "int1_t".
                        filter = createBigintValuesFilter<bool>(elements, offset, size);
                        break;
                    case TypeKind::TIMESTAMP:
                        filter = createBigintValuesFilter<Timestamp>(elements, offset, size);
                        break;
                    case TypeKind::VARCHAR:
                    case TypeKind::VARBINARY:
                        filter = createBytesValuesFilter(elements, offset, size);
                        break;
                    case TypeKind::UNKNOWN:
                        filter = {nullptr, true};
                        break;
                    case TypeKind::ARRAY:
                        [[fallthrough]];
                    case TypeKind::MAP:
                        [[fallthrough]];
                    case TypeKind::ROW:
                        return VectorSetInPredicate::create(elements, offset, size);
                    default:
                        POLLUX_UNSUPPORTED(
                            "Unsupported in-list type for IN predicate: {}",
                            inListType->toString());
                }
                return std::make_shared<InPredicate>(
                    std::move(filter.first), filter.second);
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                if (alwaysNull_) {
                    auto localResult = createBoolConstantNull(rows.end(), context);
                    context.moveOrCopyResult(localResult, rows, result);
                    return;
                }

                const auto &input = args[0];
                switch (input->type_kind()) {
                    case TypeKind::HUGEINT:
                        applyTyped<int128_t>(rows, input, context, result, [&](int128_t value) {
                            return filter_->testInt128(value);
                        });
                        break;
                    case TypeKind::BIGINT:
                        applyTyped<int64_t>(rows, input, context, result, [&](int64_t value) {
                            return filter_->testInt64(value);
                        });
                        break;
                    case TypeKind::INTEGER:
                        applyTyped<int32_t>(rows, input, context, result, [&](int32_t value) {
                            return filter_->testInt64(value);
                        });
                        break;
                    case TypeKind::SMALLINT:
                        applyTyped<int16_t>(rows, input, context, result, [&](int16_t value) {
                            return filter_->testInt64(value);
                        });
                        break;
                    case TypeKind::TINYINT:
                        applyTyped<int8_t>(rows, input, context, result, [&](int8_t value) {
                            return filter_->testInt64(value);
                        });
                        break;
                    case TypeKind::REAL:
                        applyTyped<float>(rows, input, context, result, [&](float value) {
                            if (value == float{}) {
                                value = 0;
                            } else if (std::isnan(value)) {
                                // We de-normalize NaN values to ensure different binary
                                // representations
                                // are treated the same.
                                value = std::numeric_limits<float>::quiet_NaN();
                            }
                            return filter_->testInt64(reinterpret_cast<const int32_t &>(value));
                        });
                        break;
                    case TypeKind::DOUBLE:
                        applyTyped<double>(rows, input, context, result, [&](double value) {
                            if (value == double{}) {
                                value = 0;
                            } else if (std::isnan(value)) {
                                // We de-normalize NaN values to ensure different binary
                                // representations
                                // are treated the same.
                                value = std::numeric_limits<double>::quiet_NaN();
                            }
                            return filter_->testInt64(reinterpret_cast<const int64_t &>(value));
                        });
                        break;
                    case TypeKind::BOOLEAN:
                        applyTyped<bool>(rows, input, context, result, [&](bool value) {
                            return filter_->testInt64(value);
                        });
                        break;
                    case TypeKind::TIMESTAMP:
                        applyTyped<Timestamp>(
                            rows, input, context, result, [&](Timestamp value) {
                                return filter_->testInt64(value.toMillis());
                            });
                        break;
                    case TypeKind::VARCHAR:
                    case TypeKind::VARBINARY:
                        applyTyped<StringView>(
                            rows, input, context, result, [&](StringView value) {
                                return filter_->testBytes(value.data(), value.size());
                            });
                        break;
                    default:
                        POLLUX_UNSUPPORTED(
                            "Unsupported input type for the IN predicate: {}",
                            input->type()->toString());
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    // (T, array(T)) -> boolean for constant IN lists.
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("boolean")
                    .argumentType("T")
                    .constantArgumentType("array(T)")
                    .build(),
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("boolean")
                    .argumentType("T")
                    .constantArgumentType("array(unknown)")
                    .build(),
                };
            }

        private:
            static VectorPtr createBoolConstantNull(
                vector_size_t size,
                exec::EvalCtx &context) {
                return std::make_shared<ConstantVector<bool> >(
                    context.pool(), size, true /*isNull*/, BOOLEAN(), false);
            }

            static VectorPtr
            createBoolConstant(bool value, vector_size_t size, exec::EvalCtx &context) {
                return std::make_shared<ConstantVector<bool> >(
                    context.pool(), size, false /*isNull*/, BOOLEAN(), std::move(value));
            }

            template<typename T, typename F>
            void applyTyped(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                exec::EvalCtx &context,
                VectorPtr &result,
                F &&testFunction) const {
                POLLUX_CHECK(filter_, "IN predicate supports only constant IN list");

                // Indicates whether result can be true or null only, e.g. no false results.
                const bool passOrNull = filter_->testNull();

                if (arg->is_constant_encoding()) {
                    auto simpleArg = arg->as_unchecked<SimpleVector<T> >();
                    VectorPtr localResult;
                    if (simpleArg->is_null_at(rows.begin())) {
                        localResult = createBoolConstantNull(rows.end(), context);
                    } else {
                        static const SelectivityVector oneRow(1, true);
                        context.applyToSelectedNoThrow(oneRow, [&](vector_size_t row) {
                            bool pass = testFunction(simpleArg->value_at(rows.begin()));
                            if (!pass && passOrNull) {
                                localResult = createBoolConstantNull(rows.end(), context);
                            } else {
                                localResult = createBoolConstant(pass, rows.end(), context);
                            }
                        });
                    }

                    context.moveOrCopyResult(localResult, rows, result);
                    return;
                }

                POLLUX_CHECK_EQ(arg->encoding(), VectorEncoding::Simple::FLAT);
                auto flatArg = arg->as_unchecked<FlatVector<T> >();

                context.ensure_writable(rows, BOOLEAN(), result);
                result->clear_nulls(rows);
                auto *boolResult = result->as_unchecked<FlatVector<bool> >();

                auto *rawResults = boolResult->mutableRawValues<uint64_t>();

                if (flatArg->may_have_nulls() || passOrNull) {
                    context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
                        if (flatArg->is_null_at(row)) {
                            boolResult->set_null(row, true);
                        } else {
                            bool pass = testFunction(flatArg->valueAtFast(row));
                            if (!pass && passOrNull) {
                                boolResult->set_null(row, true);
                            } else {
                                bits::setBit(rawResults, row, pass);
                            }
                        }
                    });
                } else {
                    context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
                        bool pass = testFunction(flatArg->valueAtFast(row));
                        bits::setBit(rawResults, row, pass);
                    });
                }
            }

            const std::unique_ptr<common::Filter> filter_;
            const bool alwaysNull_;
        };
    } // namespace

    POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION(
        udf_in,
        InPredicate::signatures(),
        InPredicate::create);
} // namespace kumo::pollux::functions
